import os
from moviepy.video.io.VideoFileClip import VideoFileClip
import numpy as np
from PIL import Image

def extract_frames(video_file, output_dir, fps, quality):
    """
    从视频文件中提取帧并将其保存为JPEG图像。

    :param video_file: 视频文件路径
    :param output_dir: 输出文件夹路径
    :param fps: 帧率（每秒帧数）
    :param quality: 输出JPEG的质量
    """

    # 获取视频文件名作为照片前缀
    file_name = os.path.splitext(os.path.basename(video_file))[0]

    # 创建输出文件夹
    output_path = os.path.join(output_dir, file_name)
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # 打开视频文件
    clip = VideoFileClip(video_file)

    # 如果设置的帧率大于视频的帧率，则使用视频的帧率
    if fps > clip.fps:
        fps = clip.fps

    # 计算每一帧之间的时间间隔
    time_interval = 1.0 / fps

    # 逐帧提取并保存照片
    for t in range(int(clip.duration / time_interval)):
        frame = clip.get_frame(t * time_interval)
        pil_image = Image.fromarray(frame)
        file_path = os.path.join(output_path, f"{file_name}_{t}.jpg")
        pil_image.save(file_path, quality=quality)

def extract_frames_from_folder(folder_path, output_dir, fps, quality):
    """
    从指定文件夹中的所有视频文件中提取帧并将其保存为JPEG图像。

    :param folder_path: 视频文件所在文件夹路径
    :param output_dir: 输出文件夹路径
    :param fps: 帧率（每秒帧数）
    :param quality: 输出JPEG的质量
    """

    # 遍历文件夹下的所有文件
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            # 仅处理扩展名为 .mp4 或 .avi 的文件
            if file.endswith(".mp4") or file.endswith(".avi"):
                file_path = os.path.join(root, file)
                extract_frames(file_path, output_dir, fps, quality)

if __name__ == '__main__':
    # 输入参数
    input_folder = "./mp4/20230618_wangjian"
    output_folder = "./jpg/20230618_wangjian"
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    fps = 30
    quality = 100

    # 处理所有视频文件
    print('extract_frames_from_folder is going......')
    extract_frames_from_folder(input_folder, output_folder, fps, quality)
    print('extract_frames_from_folder is Done!')